import os
import random
import shutil

# 配置路径
dataset_root = "./data/ECSSD"
source_image_dir = os.path.join(dataset_root, "Imgs")
source_mask_dir = os.path.join(dataset_root, "mask")
train_root = os.path.join(dataset_root, "train")
test_root = os.path.join(dataset_root, "test")

# 创建目标文件夹
os.makedirs(os.path.join(train_root, "Imgs"), exist_ok=True)
os.makedirs(os.path.join(train_root, "mask"), exist_ok=True)
os.makedirs(os.path.join(test_root, "Imgs"), exist_ok=True)
os.makedirs(os.path.join(test_root, "mask"), exist_ok=True)

# 获取所有文件名（不带扩展名）
all_files = [f.split(".")[0] for f in os.listdir(source_image_dir) if f.endswith(('.jpg', '.png'))]
random.shuffle(all_files)  # 随机打乱

# 划分数据集 (70% 训练, 30% 测试)
split_idx = int(len(all_files) * 0.7)
train_files = all_files[:split_idx]  # 700 张
test_files = all_files[split_idx:]  # 300 张

# 创建训练集
for file in train_files:
    # 复制图像
    src_img = os.path.join(source_image_dir, f"{file}.jpg")
    dst_img = os.path.join(train_root, "Imgs", f"{file}.jpg")
    shutil.copy(src_img, dst_img)

    # 复制掩码
    src_mask = os.path.join(source_mask_dir, f"{file}.png")
    dst_mask = os.path.join(train_root, "mask", f"{file}.png")
    shutil.copy(src_mask, dst_mask)

# 创建测试集
for file in test_files:
    src_img = os.path.join(source_image_dir, f"{file}.jpg")
    dst_img = os.path.join(test_root, "Imgs", f"{file}.jpg")
    shutil.copy(src_img, dst_img)

    src_mask = os.path.join(source_mask_dir, f"{file}.png")
    dst_mask = os.path.join(test_root, "mask", f"{file}.png")
    shutil.copy(src_mask, dst_mask)

print(f"创建训练集: {len(train_files)} 张图片")
print(f"创建测试集: {len(test_files)} 张图片")